#! /usr/bin/python
#
# Stephen Poletto (spoletto)
# Peter Wilmot (pbwilmot)
# CSCI1580 - Web Search
# Spring 2011 - Brown University
#
# Reads in a list of queries
# and returns matching documents
# using the provided index file.

from bool_parser import bool_expr_ast
from positionalIndex import *
from porterStemmer import *
import marshal
import string
import sys
import re

if (len(sys.argv) != 4):
    print ""
    print "usage: queryIndex <stopwords> <index> <titleindex>"
    print ""
    sys.exit()

# Read the index data into memory
serializedFile = open(sys.argv[2])
index = PositionalIndex(marshal.load(serializedFile))
serializedFile.close()

# Read in the stop words
stopWordsFile = open(sys.argv[1], "r")
stopWords = []
for stop in stopWordsFile.readlines():
    word = stop.rstrip("\n")
    stopWords.append(word)

# Build the Porter Stemmer
stemmer = PorterStemmer()

# Print the results to stdout
def print_results(setOfDocIDs):
    successfulDocuments = list(setOfDocIDs)
    successfulDocuments.sort(key=int)
    if len(successfulDocuments) > 0:
        for id in successfulDocuments[0:-1]:
            sys.stdout.write(id + " ")
        sys.stdout.write(successfulDocuments[-1])
    sys.stdout.write('\n')

# For processing one word queries.
# Simply look up the word in the
# dictionary, and print the postings
# list of document IDs.

def processOWQ_private(query):
    docIDs = []
    docIDToPositions = index.lookup(query)
    if docIDToPositions != None:
        docIDs = docIDToPositions.keys()
    return set(docIDs)

def processOWQ(query):
    print_results(processOWQ_private(query))

# For processing free text queries.
# Simply look up each word in the
# dictionary, and print the union
# of documents from the associated
# postings lists.

def processFTQ(keywordList):
    resultDocs = set([])
    for word in keywordList:
        docIDToPositions = index.lookup(word)
        if docIDToPositions != None:
            docIDs = docIDToPositions.keys()
            resultDocs = resultDocs.union(set(docIDs))
    print_results(resultDocs)

# For processing phrase queries.
# Each term in the keyword list
# must occur within a single
# document, so first we take
# the intersection of documents
# from the postings lists associated
# with each term. Then, once we have
# our candidate documents, we check
# the positional indices to make sure
# the phrase appears as specified
# in the query.

def processPQ(keywordList):
    candidateDocs = None
    wordToDocList = {}
    for word in keywordList:
        docIDToPositions = index.lookup(word)
        if docIDToPositions != None:
            # Cache the postings list for each word.
            wordToDocList[word] = docIDToPositions
            docIDs = docIDToPositions.keys()
            # Initialize the set of candidate docs.
            if candidateDocs == None:
                candidateDocs = set(docIDs)
            candidateDocs = candidateDocs.intersection(set(docIDs))
        else:
            sys.stdout.write('\n')
            return

    # Candidate documents contain every word in the phrase.
    successfulDocuments = []
    for document in candidateDocs:
        positionsOfWordPrior = None
        for word in keywordList:
            positionsOfWord = wordToDocList[word][document]
            # First word in the query
            if positionsOfWordPrior == None:
                positionsOfWordPrior = positionsOfWord
            else:
                newPositionsOfWordPrior = []
                for position in positionsOfWord:
                    if position-1 in positionsOfWordPrior:
                        newPositionsOfWordPrior.append(position)
                positionsOfWordPrior = newPositionsOfWordPrior
                if len(positionsOfWordPrior) == 0:
                    break
        # If we made it through the full phrase with the last
        # term succesfully matched up with a prior term, we
        # know the document contained the entire phrase.
        if len(positionsOfWordPrior) > 0:
            successfulDocuments.append(document)
    
    # Print results
    print_results(successfulDocuments)

# For processing boolean queries.
# The AST generated by the TA
# support code will be passed in.
# The first element of the AST is
# always the operation. The subsequent
# items may be stings (single word
# base case) or tuples, which require
# recursive evaluation.

def processBQ_private(boolean_ast):
    results = None
    operation = boolean_ast[0]
    for operand in boolean_ast[1]:
        matches = set([])
        # Tuples require further parsing
        if isinstance(operand, tuple):
            matches = processBQ_private(operand)
        else:
            stemmed = stemmedQuery(operand)
            if not len(stemmed) < 1:
                # Word is not a stop word.
                matches = processOWQ_private(stemmed[0])
        if results == None:
            results = matches
        if operation == "OR":
            results = results.union(matches)
        elif operation == "AND":
            results = results.intersection(matches)
    return results

def processBQ(boolean_ast):
    results = processBQ_private(boolean_ast)
    print_results(results)

def removeOperatorsFromString(string, operators):
    for op in operators:
        string = string.replace(op, "")
    return string

def stemmedQuery(query):
    stemmed = removeOperatorsFromString(query, ["AND", "OR", "(", ")", "\""])
    stemmed = stemmed.lower()
    stemmed = re.compile(r'\b[a-z0-9]+\b').findall(stemmed)
    outputStream = []
    for word in stemmed:
        isStopWord = False;
        for stopWord in stopWords:
            if word == stopWord:
                isStopWord = True
        # If the word is in stopwords, ignore it.
        if not isStopWord:
            outputStream.append(stemmer.stem(word, 0, len(word)-1))
    return outputStream

def processQuery(query):
    if len(stemmedQuery(query)) < 1:
        # Search query has zero terms!
        sys.stdout.write("\n")
        return
    if len(re.compile(r'\b\w+\b').findall(query)) == 1:
        # Single word query:
        processOWQ(stemmedQuery(query)[0])
    elif query.find("\"") != -1:
        processPQ(stemmedQuery(query))
    elif query != bool_expr_ast(query):
        # We have a bool AST; must be a BQ
        processBQ(bool_expr_ast(query))
    else:
        processFTQ(stemmedQuery(query))

# Main run loop.
# Read input from stdin and process
# each incoming query.
while True:
    try:
        query = raw_input()
        processQuery(query.rstrip('\n'))
    except EOFError:
        break